using System;
using System.Collections;
using System.Collections.Generic;
using UnityEngine;
using WeightAndBias = JufGame.WeightBiasMemory.LayerWeightAndBias;
using Random = UnityEngine.Random;

namespace JufGame
{
	public class GA : MonoBehaviour
	{		
		[Serializable]
		private class GAOne //遗传算法（genetic algorithm，简写GA）中的个体
		{
			public float fitness; //适应度
			public WeightAndBias[] WB; //网络的权重和偏置集合
			public GAOne(WeightAndBias[] source)
			{
				fitness = 0;
				WB = new WeightAndBias[source.Length];
				for(int i = 0; i < source.Length; ++i)
				{
					WB[i].inputCount = source[i].inputCount;
					WB[i].outputCount = source[i].outputCount;
					WB[i].weights = new float[source[i].weights.Length];
					WB[i].bias = new float[source[i].bias.Length];
				}
				WeightBiasMemory.RandomWeightAndBias(ref WB, -1, 1, true);
			}
		}

		[Header("遗传算法设置")]
		public int AllPopulation = 50; //总人口数量设置
		[Range(0, 1)]
		public float mutationRate = 0.3f; // 突变率
		public float mutationScale = 1; // 突变幅度
		public GAUnit TrainUnit; //正在训练的具体个体

		[Header("当前个体信息")]
		public int curIndex;
		private GAOne[] parents; //父代
		private GAOne[] children; //子代
		private GAOne[] totalPopulation; //父代+子代，用于排序，优胜劣汰
		[SerializeField] private bool isTrainChildren; //当前是否正在训练子代
		[SerializeField] private bool isEndTrain; //是否训练结束
		private float[] cumulativeProbabilities; //用于轮盘赌的积累概率

		private void Start()
		{
			parents = new GAOne[AllPopulation];
			children = new GAOne[AllPopulation];
			totalPopulation = new GAOne[AllPopulation * 2];
			cumulativeProbabilities = new float[AllPopulation];
			for(int i = 0; i < AllPopulation; ++i)
			{
				parents[i] = new GAOne(TrainUnit.memory.WeiBiasArray);
				children[i] = new GAOne(TrainUnit.memory.WeiBiasArray);
			}
			curIndex = 0;
			WeightBiasMemory.DeepCopyAllLayerWB(ref parents[0].WB, ref TrainUnit.memory.WeiBiasArray);
		}
		
		private void FixedUpdate()
		{
			if(isEndTrain)
			{
				SaveBest();
			}
			else if(TrainUnit.isOver)
			{
				parents[curIndex].fitness = TrainUnit.FitNess;
				TrainUnit.ReStart();
				if(++curIndex < AllPopulation)
				{
					WeightBiasMemory.DeepCopyAllLayerWB(ref parents[curIndex].WB, ref TrainUnit.memory.WeiBiasArray);
				}
				else if(!isTrainChildren)
				{
					isTrainChildren = true;
					CalcRouletteWheel();
					GetChild();
					for(int i = 0; i < AllPopulation; ++i)
					{
						children[i].fitness = parents[i].fitness;
						WeightBiasMemory.DeepSwap(ref children[i].WB, ref parents[i].WB);
					}
					curIndex = 0;
					WeightBiasMemory.DeepCopyAllLayerWB(ref parents[0].WB, ref TrainUnit.memory.WeiBiasArray);
				}
				else
				{
					isTrainChildren = false;
					GetBest();
					for(int i = 0; i < AllPopulation; ++i)
					{
						WeightBiasMemory.DeepCopyAllLayerWB(ref totalPopulation[i].WB, ref parents[i].WB);
					}
					curIndex = 0;
					WeightBiasMemory.DeepCopyAllLayerWB(ref parents[0].WB, ref TrainUnit.memory.WeiBiasArray);
				}
			}
		}

		//在父代和子代组成的整体中选出适应度高的新父代
		private void GetBest()
		{
			for(int i = 0; i < totalPopulation.Length; ++i)
			{
				if (i < AllPopulation)
					totalPopulation[i] = parents[i];
				else
					totalPopulation[i] = children[i - AllPopulation];
			}
			Array.Sort(totalPopulation, (a, b) => b.fitness.CompareTo(a.fitness));
		}

		//计算轮盘赌概率分布
		private void CalcRouletteWheel()
		{
			float totalFitness = 0f;
			for (int i = 0; i < parents.Length; i++)
			{
				totalFitness += parents[i].fitness;
			}
			float cumulativeSum = 0f;
			for (int i = 0; i < cumulativeProbabilities.Length; i++)
			{
				cumulativeSum += (parents[i].fitness / totalFitness);
				cumulativeProbabilities[i] = cumulativeSum;
			}
		}

		//轮盘赌随机下标
		private int GetRouletteRandom()
		{
			float rand = Random.value;
			// 选择个体
			for (int i = 0; i < cumulativeProbabilities.Length; i++)
			{
				if (rand < cumulativeProbabilities[i])
				{
					return i;
				}
			}
			// 如果没有找到，返回最后一个个体（通常不会发生）
			return cumulativeProbabilities.Length - 1;
		}

		private void GetChild()
		{
			int p1, p2;
			for(int i = 0; i < parents.Length; i += 2)
			{
				p2 = p1 = GetRouletteRandom();
				var curWB = parents[i].WB;
				while(p1 == p2 && parents.Length > 1)
				{
					p2 = GetRouletteRandom();
				}
				for(int j = 0; j < curWB.Length; ++j)
				{
					var curW = curWB[j].weights;
					for (int k = 0; k < curW.Length; ++k)
					{
						if(Random.value < 0.5)
						{
							children[i].WB[j].weights[k] = parents[p2].WB[j].weights[k];
							if (i + 1 < children.Length)
							{
								children[i + 1].WB[j].weights[k] = parents[p1].WB[j].weights[k];
							}
						}
						else
						{
							children[i].WB[j].weights[k] = parents[p1].WB[j].weights[k];
							if (i + 1 < children.Length)
							{
								children[i + 1].WB[j].weights[k] = parents[p2].WB[j].weights[k];
							}
						}
						if (Random.value < mutationRate) //随机变异，mutationRate为变异率
						{
							//mutationScale为变异的幅度，即变异带来的数值增减幅度
							children[i].WB[j].weights[k] += Random.Range(-mutationScale, mutationScale);
						}
						if (i + 1 < children.Length && Random.value < mutationRate)
						{
							children[i + 1].WB[j].weights[k] += Random.Range(-mutationScale, mutationScale);
						}
					}
					var curB = curWB[j].bias;
					for (int k = 0; k < curB.Length; ++k)
					{
						if(Random.value < 0.5)
						{
							children[i].WB[j].bias[k] = parents[p2].WB[j].bias[k];
							if (i + 1 < children.Length)
							{
								children[i + 1].WB[j].bias[k] = parents[p1].WB[j].bias[k];
							}
						}
						else
						{
							children[i].WB[j].bias[k] = parents[p1].WB[j].bias[k];
							if (i + 1 < children.Length)
							{
								children[i + 1].WB[j].bias[k] = parents[p2].WB[j].bias[k];
							}
						}
						if (Random.value < mutationRate) //随机变异，mutationRate为变异率
						{
							//mutationScale为变异的幅度，即变异带来的数值增减幅度
							children[i].WB[j].bias[k] += Random.Range(-mutationScale, mutationScale);
						}
						if (i + 1 < children.Length && Random.value < mutationRate)
						{
							children[i + 1].WB[j].bias[k] += Random.Range(-mutationScale, mutationScale);
						}
					}
				}
			}
		}

		private void SaveBest()
		{
			var findArr = isTrainChildren ? children : parents;
			WeightBiasMemory.DeepCopyAllLayerWB(ref findArr[0].WB, ref TrainUnit.memory.WeiBiasArray);
		}
	}
}
